import pandas as pd
import numpy as np
import logging
import time
from datetime import datetime, time as dt_time
from threading import Thread, Lock
from queue import Queue

class AutoTrader:
    def __init__(self, 
                 broker_api,  # 券商API接口
                 strategy,    # 交易策略
                 symbols,     # 交易标的列表
                 initial_capital=100000,
                 max_positions=5,
                 position_size=0.2,
                 stop_loss=0.02,
                 take_profit=0.05):
        """
        初始化自动交易系统
        
        参数:
            broker_api: 券商API接口实例
            strategy: 交易策略实例
            symbols: 交易标的列表

            initial_capital: 初始资金
            max_positions: 最大持仓数量
            position_size: 每个持仓的资金比例
            stop_loss: 止损比例
            take_profit: 止盈比例
        """
        self.broker = broker_api
        self.strategy = strategy
        self.symbols = symbols
        self.capital = initial_capital
        self.max_positions = max_positions
        self.position_size = position_size
        self.stop_loss = stop_loss
        self.take_profit = take_profit
        
        # 初始化数据结构
        self.positions = {}  # 当前持仓
        self.orders = {}     # 待执行订单
        self.order_queue = Queue()  # 订单队列
        
        # 线程锁
        self.position_lock = Lock()
        self.order_lock = Lock()
        
        # 初始化日志
        self._setup_logger()
        
        # 市场状态
        self.is_trading = False
        
    def _setup_logger(self):
        """设置日志"""
        self.logger = logging.getLogger('AutoTrader')
        self.logger.setLevel(logging.INFO)
        
        # 创建文件处理器
        fh = logging.FileHandler('auto_trader.log')
        fh.setLevel(logging.INFO)
        
        # 创建控制台处理器
        ch = logging.StreamHandler()
        ch.setLevel(logging.INFO)
        
        # 创建格式器
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)
        
        # 添加处理器
        self.logger.addHandler(fh)
        self.logger.addHandler(ch)

    def start(self):
        """启动自动交易"""
        self.is_trading = True
        self.logger.info("启动自动交易系统")
        
        # 启动线程
        Thread(target=self._market_data_thread).start()
        Thread(target=self._order_processing_thread).start()
        Thread(target=self._risk_management_thread).start()

    def stop(self):
        """停止自动交易"""
        self.is_trading = False
        self.logger.info("停止自动交易系统")

    def _market_data_thread(self):
        """市场数据处理线程"""
        while self.is_trading:
            try:
                # 检查是否在交易时间
                if not self._is_trading_time():
                    time.sleep(60)
                    continue
                
                # 获取市场数据
                for symbol in self.symbols:
                    market_data = self.broker.get_market_data(symbol)
                    if market_data is None:
                        continue
                    
                    # 运行策略
                    signal = self.strategy.generate_signal(market_data)
                    
                    # 处理交易信号
                    if signal != 0:
                        self._handle_signal(symbol, signal, market_data)
                        
                time.sleep(1)  # 避免请求过于频繁
                
            except Exception as e:
                self.logger.error(f"市场数据处理错误: {str(e)}")
                time.sleep(5)

    def _order_processing_thread(self):
        """订单处理线程"""
        while self.is_trading:
            try:
                if not self.order_queue.empty():
                    order = self.order_queue.get()
                    self._execute_order(order)
                time.sleep(0.1)
                
            except Exception as e:
                self.logger.error(f"订单处理错误: {str(e)}")
                time.sleep(1)

    def _risk_management_thread(self):
        """风险管理线程"""
        while self.is_trading:
            try:
                with self.position_lock:
                    for symbol, position in self.positions.items():
                        current_price = self.broker.get_latest_price(symbol)
                        if current_price is None:
                            continue
                            
                        # 检查止损
                        if self._check_stop_loss(symbol, position, current_price):
                            self._place_order(symbol, "SELL", position['volume'])
                            
                        # 检查止盈
                        elif self._check_take_profit(symbol, position, current_price):
                            self._place_order(symbol, "SELL", position['volume'])
                            
                time.sleep(1)
                
            except Exception as e:
                self.logger.error(f"风险管理错误: {str(e)}")
                time.sleep(5)

    def _handle_signal(self, symbol, signal, market_data):
        """处理交易信号"""
        try:
            with self.position_lock:
                # 买入信号
                if signal > 0 and symbol not in self.positions:
                    if len(self.positions) < self.max_positions:
                        volume = self._calculate_position_size(market_data['price'])
                        self._place_order(symbol, "BUY", volume)
                
                # 卖出信号
                elif signal < 0 and symbol in self.positions:
                    volume = self.positions[symbol]['volume']
                    self._place_order(symbol, "SELL", volume)
                    
        except Exception as e:
            self.logger.error(f"处理信号错误: {str(e)}")

    def _execute_order(self, order):
        """执行订单"""
        try:
            result = self.broker.place_order(
                symbol=order['symbol'],
                direction=order['direction'],
                volume=order['volume'],
                price=order['price']
            )
            
            if result['status'] == 'SUCCESS':
                self._update_positions(order)
                self.logger.info(f"订单执行成功: {order}")
            else:
                self.logger.warning(f"订单执行失败: {result['message']}")
                
        except Exception as e:
            self.logger.error(f"执行订单错误: {str(e)}")

    def _calculate_position_size(self, price):
        """计算持仓数量"""
        available_capital = self.capital * self.position_size
        return int(available_capital / price / 100) * 100  # 确保为100的整数倍

    def _is_trading_time(self):
        """检查是否在交易时间内"""
        now = datetime.now().time()
        morning_start = dt_time(9, 30)
        morning_end = dt_time(11, 30)
        afternoon_start = dt_time(13, 0)
        afternoon_end = dt_time(15, 0)
        
        return ((morning_start <= now <= morning_end) or 
                (afternoon_start <= now <= afternoon_end))

    def _check_stop_loss(self, symbol, position, current_price):
        """检查止损条件"""
        return ((current_price - position['price']) / position['price']) <= -self.stop_loss

    def _check_take_profit(self, symbol, position, current_price):
        """检查止盈条件"""
        return ((current_price - position['price']) / position['price']) >= self.take_profit

    def _place_order(self, symbol, direction, volume, price=None):
        """下单"""
        order = {
            'symbol': symbol,
            'direction': direction,
            'volume': volume,
            'price': price or self.broker.get_latest_price(symbol),
            'timestamp': datetime.now()
        }
        self.order_queue.put(order)
        self.logger.info(f"下单: {order}")

    def _update_positions(self, order):
        """更新持仓"""
        with self.position_lock:
            if order['direction'] == "BUY":
                self.positions[order['symbol']] = {
                    'volume': order['volume'],
                    'price': order['price'],
                    'timestamp': order['timestamp']
                }
            else:  # SELL
                if order['symbol'] in self.positions:
                    del self.positions[order['symbol']]
